package com.andfo.testframework;

import java.util.Properties;
import java.io.*;
import java.sql.DriverManager;
import java.sql.Connection;

import com.mysql.jdbc.jdbc2.optional.MysqlConnectionPoolDataSource;
import com.mysql.jdbc.jdbc2.optional.MysqlDataSource;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.ext.mysql.MySqlConnection;

import javax.sql.ConnectionPoolDataSource;
import javax.sql.DataSource;

/**
 * User: andfo
 * Date: Jul 14, 2009
 * <p/>
 * This class uses the MySql binary clients to backup and restore files. This is
 * significantly faster than running jdbc together with DBUnit.
 * <p/>
 * Suported properties are:
 * <ul>
 * <li>databaseName</li>
 * <li>restoreFilesPath</li>
 * <li>dbUser</li>
 * <li>dbPassword</li>
 * <li>mysqlBinaryPath</li>
 * <li>dbHost</li>
 * </ul>
 */
public class MySqlDataStore implements DataStore {

    private String databaseName;
    private String restoreFilesPath;
    private String dbUser;
    private String dbPassword;
    private String mysqlBinaryPath;
    private String dbHost;
    private int dbPort;
    private String jdbcDriver;
    private boolean remote;
    private String dbHostPassword;
    private String dbHostUser;
    private MysqlDataSource dataSource;

    private SshExec ssh;

    Log log = LogFactory.getLog(getClass());

    public void init(Properties properties) {
        databaseName = properties.getProperty("databaseName", "helios_test");
        restoreFilesPath = properties.getProperty("restoreFilesPath", "./restorefiles");
        dbHost = properties.getProperty("dbHost", "localhost");
        dbUser = properties.getProperty("dbUser", "root");
        dbPassword = properties.getProperty("dbPassword", "");
        mysqlBinaryPath = properties.getProperty("mysqlBinaryPath", "/usr/bin");
        dbPort = Integer.parseInt(properties.getProperty("dbPort", "3306"));
        jdbcDriver = properties.getProperty("jdbcDriver", "com.mysql.jdbc.Driver");
        remote = properties.getProperty("mysqlRemote", "false").equals("true");
        dbHostPassword = properties.getProperty("dbHostPassword");    // default null
        dbHostUser = properties.getProperty("dbHostUser", "root");

        if(remote) {
            ssh = SshExec.getConnection(dbHostUser, dbHostPassword, dbHost);
        }
    }

    /**
     * This will not work over remote connections
     * */
    public IDatabaseConnection createConnection() throws Exception {
        try {
            return new MySqlConnection(getDataSource().getConnection(), databaseName);
        } catch (Exception e) {
            log.error("Cannot connect to the database");
            throw e;
        }
    }

    public synchronized DataSource getDataSource() {
        if(dataSource == null) {
            dataSource = new MysqlDataSource();
            dataSource.setUser(dbUser);
            dataSource.setPassword(dbPassword);
            dataSource.setServerName(dbHost);
            dataSource.setPort(dbPort);
            dataSource.setDatabaseName(databaseName);
        }
        return dataSource;
    }


    public void loadRestoreFile(String restoreName) {

        File restoreFile = new File(restoreFilesPath + "/" + restoreName + ".sql");

        int total = 0;
        long started = System.currentTimeMillis();
        try {

            String command = String.format("%s/mysql -h %s --port=%s -u %s %s " + (dbPassword.length() == 0 ? "" : "-p" + dbPassword), mysqlBinaryPath, dbHost, dbPort, dbUser, databaseName);
            if (remote) {
                // Construct the file name. If we get it from the File object it might replace / with \.
                String restoreFileName = restoreFilesPath + "/" + restoreName + ".sql";
                ssh.exec(command + " < " + restoreFileName);
            } else {

                if (!restoreFile.exists())
                    throw new FileNotFoundException("Restore file does not exist:" + restoreFile.getAbsolutePath());


                log.info("Running command:" + command);
                Process p = Runtime.getRuntime().exec(command);

                // Read and log process outputs
                AsyncStreamReader.startReader(p.getInputStream(), "STDOUT");
                AsyncStreamReader.startReader(p.getErrorStream(), "ERROR");

                log.info("Restoring from file:" + restoreFile);

                BufferedInputStream bis = new BufferedInputStream(new FileInputStream(restoreFile));
                OutputStream out = new BufferedOutputStream(p.getOutputStream());
                byte[] buf = new byte[4096];
                int read;
                while ((read = bis.read(buf)) != -1) {
                    total += read;
                    out.write(buf, 0, read);
                }

                out.close();
                bis.close();

                p.waitFor();

                log.info("Restore finished. Total bytes read: " + total + ", time: " + ((System.currentTimeMillis() - started) / 1000) + "s");
            }
        } catch (FileNotFoundException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }

    }

    public void createRestoreFile(String restoreName) {
        int total = 0;
        long started = System.currentTimeMillis();

        try {
            String restoreFile = restoreFilesPath + "/" + restoreName + ".sql";
            String command = String.format("%s/mysqldump -h %s --port=%s -u %s --opt " + (dbPassword.length() == 0 ? "" : "-p" + dbPassword) + " %s", mysqlBinaryPath, dbHost, dbPort, dbUser, databaseName);

            if (remote) {
                ssh.exec(command + " > " + restoreFile);
            } else {
                log.info("Running command:" + command);
                Process p = Runtime.getRuntime().exec(command);

                // Read and log error stream
                AsyncStreamReader.startReader(p.getErrorStream(), "ERROR");

                log.info("Creating restore file: " + restoreFile);

                InputStream bis = p.getInputStream();
                BufferedOutputStream bos = new BufferedOutputStream(new FileOutputStream(restoreFile));
                byte[] buf = new byte[4096];
                int read;
                while ((read = bis.read(buf)) != -1) {
                    total += read;
                    bos.write(buf, 0, read);
                }


                bos.close();
                bis.close();

                if (p.waitFor() != 0) {
                    throw new RuntimeException("Command:" + command + " failed.");
                }
            }

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        log.info("Total bytes written:" + total + ", time: " + ((System.currentTimeMillis() - started) / 1000) + "s");
    }

    public boolean restoreFileExists(String restoreName) {
        String restoreFile = restoreFilesPath + "/" + restoreName + ".sql";
        if(remote) {
            // If the file exists, ls will then return the exact name of the file
            return ssh.exec("ls " + restoreFile).trim().equals(restoreFile);
        } else {
            return new File(restoreFile).exists();
        }
    }

    public void recreateDb() {
        String ddl = String.format("DROP DATABASE IF EXISTS `%s`; CREATE DATABASE `%s`;", databaseName, databaseName);
        String command = String.format("%s/mysql -h %s --port=%s -u %s "  + (dbPassword.length() == 0 ? "" : "-p" + dbPassword) + " -e \"%s\"", mysqlBinaryPath, dbHost, dbPort, dbUser, ddl);
    }


}
